# test_tts.py
import torchaudio
from cosyvoice.cli.cosyvoice import CosyVoice2
from cosyvoice.utils.file_utils import load_wav

# 初始化模型
cosyvoice = CosyVoice2(
    'pretrained_models/CosyVoice2-0.5B',
    load_jit=False,
    load_trt=False,
    load_vllm=False,
    fp16=False
)

# 加载参考音频
reference_audio = load_wav("./asset/ccc.wav", 16000)

# 生成测试语音
text = "我在！"
results = list(cosyvoice.inference_zero_shot(
    text,
    "9月4日，中共中央总书记",
    reference_audio,
    stream=False,
    text_frontend=False
))

# 保存结果
if results:
    speech_data = results[0]['tts_speech']

    # 调试输出张量形状
    print(f"原始张量形状: {speech_data.shape}")

    # 修正张量维度
    if speech_data.dim() == 1:
        # 一维张量 [采样点数] -> [通道数=1, 采样点数]
        speech_data = speech_data.unsqueeze(0)
    elif speech_data.dim() == 3:
        # 三维张量 [?, ?, 采样点数] -> [通道数, 采样点数]
        # 尝试移除不必要的维度
        if speech_data.size(0) == 1:
            speech_data = speech_data.squeeze(0)
        elif speech_data.size(1) == 1:
            speech_data = speech_data.squeeze(1)
        else:
            # 如果无法直接处理，取第一个通道
            speech_data = speech_data[0]

    print(f"修正后张量形状: {speech_data.shape}")

    # 确保是2D张量
    if speech_data.dim() != 2:
        print(f"错误: 无法修正张量形状为2D, 当前维度: {speech_data.dim()}")
    else:
        torchaudio.save("test_output.wav", speech_data, cosyvoice.sample_rate)
        print("测试成功! 输出保存为 test_output.wav")
else:
    print("测试失败: 没有生成语音")